#!/usr/bin/python
"""
MP4 (IsoMedia) file format
"""
import sys
sys.path.append('/home/noam/programs/construct/svn')
from construct import *


class PrintContext(Construct):
    def _parse(self, stream, context):
        print '--> box'
        print context
        print '-----'
    def _sizeof(self, context=None):
        return 0

def fourcc(name): return String(name, 4)
    
box_header = Struct('box_header',
                    UBInt32('size'),
                    If(lambda ctx: ctx['size'] == 1,
                       UBInt64('largesize')
                       ),
                    If(lambda ctx: ctx['size'] > 0,
                       Struct('type_fields',
                              fourcc('type'),
                              If(lambda ctx: ctx['type'] == 'uuid',
                                 String('extended_type', 16),
                                 ),
                              ),
                       ),
                    )


fullbox_fields = Struct('fullbox_fields',
                        UBInt8('version'),
                        Array(3, UBInt8('flags')), # TODO use 24-bit BitStruct
                        )

# File type
ftyp_box = Struct('ftyp',
                  fourcc('major_brand'),
                  fourcc('minor_version'),
                  Anchor('before_comp_brands'),
                  Array(lambda ctx: (ctx['_']['_']['box_header']['size']
                                     - (ctx['before_comp_brands']
                                        - ctx['_']['_']['before_header'])
                                     ) / 4,
                        fourcc('compatible_brands')),
                  )

# Movie header
mvhd_box = Struct('mvhd',
                  fullbox_fields,
                  IfThenElse('version_test', lambda ctx: ctx['fullbox_fields']['version'] == 1,
                             Embedded(Struct('version1',
                                             UBInt64('creation_time'),
                                             UBInt64('modification_time'),
                                             UBInt32('timescale'),
                                             UBInt64('duration'),
                                             )),
                             Embedded(Struct('version0',
                                             UBInt32('creation_time'),
                                             UBInt32('modification_time'),
                                             UBInt32('timescale'),
                                             UBInt32('duration'),
                                             )),
                             ),
                  UBInt32('rate'), # TODO should be fixed point 16.16
                  UBInt16('volume'), # TODO should be fixed point 8.8 
                  Const(UBInt16("reserved"), 0),
                  Array(2, Const(UBInt32("reserved32x2"), 0)),
                  Array(9, UBInt32('matrix')),  # TODO (nice to have) verify that it is according to restrictions?
                  Array(6, UBInt32('pre_defined')), # TODO (nice to have) is this always 0?
                  UBInt32('next_track_ID'),
                  )
                  

# Track header
tkhd_box = Struct('tkhd',
                  fullbox_fields,
                  IfThenElse('version_test', lambda ctx: ctx['fullbox_fields']['version'] == 1,
                             Embedded(Struct('version1',
                                             UBInt64('creation_time'),
                                             UBInt64('modification_time'),
                                             UBInt32('track_ID'),
                                             Const(UBInt32('reserved0'), 0),
                                             UBInt64('duration'),
                                             )),
                             Embedded(Struct('version0',
                                             UBInt32('creation_time'),
                                             UBInt32('modification_time'),
                                             UBInt32('track_ID'),
                                             Const(UBInt32('reserved0'), 0),
                                             UBInt32('duration'),
                                             )),
                             ),
                  Array(2, Const(UBInt32("reserved32x2"), 0)),
                  UBInt16('layer'), 
                  UBInt16('alternate_group'), 
                  UBInt16('volume'),
                  Const(UBInt16("reserved"), 0),
                  Array(9, UBInt32('matrix')),  # TODO (nice to have) verify that it is according to restrictions?
                  UBInt32('width'),
                  UBInt32('height'),
                  )
                  


# Media Header Box
mdhd_box = Struct('mdhd',
                  fullbox_fields,
                  IfThenElse('version_test', lambda ctx: ctx['fullbox_fields']['version'] == 1,
                             Embedded(Struct('version1',
                                             UBInt64('creation_time'),
                                             UBInt64('modification_time'),
                                             UBInt32('timescale'),
                                             UBInt64('duration'),
                                             )),
                             Embedded(Struct('version0',
                                             UBInt32('creation_time'),
                                             UBInt32('modification_time'),
                                             UBInt32('timescale'),
                                             UBInt32('duration'),
                                             )),
                             ),
                  BitStruct('language',
                            Padding(1),
                            Bits('1', 5),
                            Bits('2', 5),
                            Bits('3', 5),
                            ),
#                  UBInt16("language"),
                  UBInt16("pre_defined"),
                  )


# Media Header Box
hdlr_box = Struct('hdlr',
                  fullbox_fields,
                  UBInt32("pre_defined"),
                  UBInt32("handler_type"),
                  Array(3, Const(UBInt32('reserved'), 0)),
                  CString('name'),
                  )
                  
# Media Header Box
stsz_box = Struct('stsz',
                  fullbox_fields,
                  UBInt32("sample_size"),
                  UBInt32("sample_count"),
                  If(lambda ctx: ctx['sample_size'] == 0,
                     Array(lambda ctx: ctx['sample_count'], UBInt32('entry_size')),
                     ),
                  )
# Media Header Box
stsc_box = Struct('stsc',
                  fullbox_fields,
                  UBInt32("entry_count"),
                  Array(lambda ctx: ctx['entry_count'],
                        Struct('entry',
                               UBInt32('first_chunk'),
                               UBInt32('samples_per_chunk'),
                               UBInt32('sample_description_index'),
                               ),
                        ),
                  )
# Media Header Box
stss_box = Struct('stss',
                  fullbox_fields,
                  UBInt32("entry_count"),
                  Array(lambda ctx: ctx['entry_count'], UBInt32('sample_number')),
                  )
                  

#-----------------------------------------------------------------------------

def remaining_data_size(ctx):
    size = (ctx['_']["box_header"]["size"]
            - (ctx['after_data'] - ctx['_']['before_header']))
#    print '--- size', size
    return size

                        
box = Struct("box",
             Anchor('before_header'),
             box_header,
#             PrintContext('a'),
             Anchor('after_header'),
             IfThenElse("optional_boxed_data", lambda ctx: ctx["box_header"]["size"] > 0,
                        Struct('boxed_data',
                               Switch("box_type", lambda ctx: ctx['_']['box_header']['type_fields']['type'],
                                      {'ftyp': ftyp_box,
                                       'moov': Pass,
                                       'mvhd': mvhd_box,
                                       'trak': Pass,
                                       'tkhd': tkhd_box,
                                       'mdia': Pass,
                                       'mdhd': mdhd_box,
                                       'hdlr': hdlr_box,
                                       'minf': Pass,
                                       'stbl': Pass,
                                       'stsz': stsz_box,
                                       'stss': stss_box,
                                       'stsc': stsc_box,
                                       },
                                      default = Value('unknown type', lambda ctx: 'unknown type'),
                                      ),
                               If(lambda ctx: ctx['box_type'] == 'unknown type',
                                  String('data', lambda ctx: (ctx['_']["box_header"]["size"]
                                                              - (ctx['_']['after_header']
                                                                 - ctx['_']['before_header']))
                                         ),
                                  ),
                               Anchor('after_data'),
                               If(lambda ctx: remaining_data_size(ctx) > 0,
                                  TunnelAdapter(String("sub_boxes", remaining_data_size),
                                                GreedyRange(LazyBound("sub_box", lambda: box)),
                                                ),
                                  ),
                               ),
                        OptionalGreedyRange(Pass),
                        ),
             )



isomedia = GreedyRepeater(box)

